import os
import time
from ultralytics import YOLO
from PIL import Image

def load_images_from_folder(folder_path):
    """加载文件夹下所有图片."""
    valid_images = [".jpg", ".jpeg", ".png", ".bmp"]
    images = []
    
    for filename in os.listdir(folder_path):
        ext = os.path.splitext(filename)[1].lower()
        if ext in valid_images:
            image_path = os.path.join(folder_path, filename)
            images.append(image_path)
    return images

def infer_images(model, image_paths):
    """对每张图片进行推理并返回平均推理时间."""
    total_time = 0
    num_images = len(image_paths)
    
    for img_path in image_paths:
        img = Image.open(img_path)  # 加载图片
        
        start_time = time.time()
        model(img)  # 执行推理（只使用CPU）
        end_time = time.time()
        
        inference_time = end_time - start_time
        total_time += inference_time
    
    # 计算平均每张图片的推理时间
    avg_inference_time = total_time / num_images if num_images > 0 else 0
    return avg_inference_time

def main():
    folder_path = './datasets/images/test'  # 设置图片文件夹路径
    model_path = './runs/detect/yolo_custom5/weights/best.pt'  # 模型路径

    # 加载模型，只使用CPU
    device = 'cpu'
    model = YOLO(model_path).to(device)

    # 加载文件夹下所有图片
    image_paths = load_images_from_folder(folder_path)
    if not image_paths:
        print(f"Error: No images found in folder '{folder_path}'")
        return

    # 进行推理并计算平均时间
    avg_time = infer_images(model, image_paths)
    print(f"Average inference time per image: {avg_time:.4f} seconds")

if __name__ == '__main__':
    main()
